import cv2 as cv
from python_ai.common.xcommon import *
import numpy as np
import os
import sys
import datetime

FILE_NAME = os.path.basename(__file__)
SAVE_DIR = os.path.join('_save', FILE_NAME)
os.makedirs(SAVE_DIR, exist_ok=True)

cap = cv.VideoCapture(0)

# make color filter
# color hue
LOW = 50
HIGH = 255
# tgt_bgr = np.uint8((31, 23, 116))  # ATTENTION unit8  # The red cup
# tgt_bgr = np.uint8((132, 118, 106))  # ATTENTION unit8  # The blue shirt
tgt_bgr = np.uint8((132, 118, 106))  # ATTENTION unit8
tgt_bgr = np.expand_dims(tgt_bgr, (0, 1))
tgt_hsv = cv.cvtColor(tgt_bgr, cv.COLOR_BGR2HSV)
print(f'tgt bgr: {tgt_bgr}, hsv: {tgt_hsv}')
# range
h_range = select_in_range(tgt_hsv[0, 0, 0], 10, 10, 0, 255)

# logo
logo_img = cv.imread('../../../../large_data/pic/opencv-logo-white.png', cv.IMREAD_COLOR)  # cv.IMREAD_COLOR for 3 channels
logo_img = imzoom(logo_img, 70)
logo_img_shape = logo_img.shape
LOGO_H = logo_img_shape[0]
LOGO_W = logo_img_shape[1]
logo_img_gray = cv.cvtColor(logo_img, cv.COLOR_BGR2GRAY)
_, mask = cv.threshold(logo_img_gray, 10, 255, cv.THRESH_BINARY)  # ATTENTION threshold returns (ret, mask)
mask_inv = cv.bitwise_not(mask)
logo_masked = cv.bitwise_and(logo_img, logo_img, mask=mask)


def empty_callback(x):
    pass


title = 'cam filter canny'
cv.namedWindow(title)
cv.createTrackbar('min', title, 255, 255, empty_callback)
cv.createTrackbar('max', title, 255, 255, empty_callback)

FPS = 25
INTERVAL = 1000 // FPS
EXIT_SET = set([ord('q'), ord('Q'), 27])
while True:
    ret, img = cap.read()

    if not ret:
        print('Cannot get image, existing ...')
        break
    shape_tuple = img.shape
    H = shape_tuple[0] - 1
    W = shape_tuple[1] - 1
    H2 = H * 2
    W2 = W * 2

    # filter color mask
    lower_border = np.uint8((h_range[0][0], LOW, LOW))  # ATTENTION unit8
    upper_border = np.uint8((h_range[0][1], HIGH, HIGH))
    hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
    mask = cv.inRange(hsv, lower_border, upper_border)
    if h_range[1] is not None:
        lower_border = np.uint8((h_range[1][0], LOW, LOW))
        upper_border = np.uint8((h_range[1][1], HIGH, HIGH))
        mask2 = cv.inRange(img, lower_border, upper_border)
        mask = cv.bitwise_or(mask, mask2)
    # show mask
    img2 = mask
    img2 = cv.cvtColor(img2, cv.COLOR_GRAY2BGR)

    # filter color
    img3 = cv.bitwise_and(img, img, mask=mask)

    # canny
    min = cv.getTrackbarPos('min', title)
    max = cv.getTrackbarPos('max', title)
    # img4 = cv.Canny(img, 100, 200)
    img4 = cv.Canny(img, min, max)
    img4 = cv.cvtColor(img4, cv.COLOR_GRAY2BGR)

    # concatenate
    img_row1 = np.concatenate((img, img2), axis=1)
    img_row2 = np.concatenate((img3, img4), axis=1)
    img_all = np.concatenate((img_row1, img_row2), axis=0)

    # separate
    cv.line(img_all, (W, 0), (W, H2), (0, 255, 0), 2)
    cv.line(img_all, (0, H), (W2, H), (0, 255, 0), 2)

    # zoom
    img_all = imzoom(img_all, 700)

    # add logo
    roi = img_all[0:LOGO_H, 0:LOGO_W]
    roi_masked = cv.bitwise_and(roi, roi, mask=mask_inv)
    mixed_roi = cv.add(logo_masked, roi_masked)
    img_all[0:LOGO_H, 0:LOGO_W] = mixed_roi

    # show
    cv.imshow(title, img_all)

    k = cv.waitKey(INTERVAL)
    if k == ord('s') or k == ord('S'):
        path = os.path.join(SAVE_DIR, ensure_filename(str(datetime.datetime.now())) + '.jpg')
        print(path)
        cv.imwrite(path, img_all)
    if k in EXIT_SET:
        break

cv.destroyAllWindows()
